{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Assembly of system with multiple domains, variables and numerics\n",
    "\n",
    "This tutorial has the dual purpose of illustrating parameter assigment in PorePy, and also showing  how to set up problems in (mixed-dimensional) geometries. It contains two examples, one covering a simple setup (the pressure equation), and a second illustrating the full generality of the coupling scheme. \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import scipy.sparse as sps\n",
    "\n",
    "import porepy as pp"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Data assignment\n",
    "We will mainly use default values for parameters, while overriding some of the values. Sets of default parameters are available for flow, transport and mechanics (elasticity). For example, initialize_default_data initializes a second order permeability tensor in the flow data, and a fourth order stiffness tensor in the mechanics data. For more details, and definitions of what the defaults are, see the modules pp/params/data.py and pp/params/parameter_dictionaries.py <br>\n",
    " The parameters are stored in a class pp.Parameters. This class is again stored in the data dictionary on each node and edge in the GridBucket, that is, in the variable d in this loop. The Paramater object can be accessed by `d[pp.PARAMETERS]`. To allow storage of parameters for several problems simultaneously (say, we want to solve a combined flow and transport problem), the Parameter class uses keywords to identify sets of parameters. This keyword must also be provided to the discretization method.\n",
    "When the parameter class is initialized with default values, the default behavior is to identify the parameters by the same keyword as is used to choose the type of default parameters (`default_parameter_type`). While this is usually good practice, we here override this behavior for illustrative purposes, using the `keyword_param_storage`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def assign_data(gb, keyword_param_storage):\n",
    "    # Method to assign data.\n",
    "    for g, d in gb:\n",
    "        # This keyword is used to define which set of default parameters to pick\n",
    "        # Replace with 'transport' or 'mechanics' if needed\n",
    "        default_parameter_type = 'flow'  \n",
    "\n",
    "        # Assign a non-default permeability, for illustrative purposes\n",
    "        if g.dim == 2:\n",
    "            kxx = 10 * np.ones(g.num_cells)\n",
    "        else:\n",
    "            kxx = 0.1 * np.ones(g.num_cells)\n",
    "\n",
    "        perm = pp.SecondOrderTensor(kxx)\n",
    "        \n",
    "        # We also set Dirichlet conditions, as the default Neumann condition\n",
    "        # gives a singular problem\n",
    "        bc = pp.BoundaryCondition(g, g.get_boundary_faces(), 'dir')\n",
    "\n",
    "        # Create a dictionary to override the default parameters\n",
    "        specified_parameters = {'second_order_tensor': perm, 'bc': bc}\n",
    "\n",
    "        # Define the \n",
    "        pp.initialize_default_data(g, d, default_parameter_type, specified_parameters,\n",
    "                                   keyword_param_storage)\n",
    "\n",
    "        # Internally to the Parameter class, the parameters are stored as dictionaries.\n",
    "        # To illustrate how to access specific sets of parameters, print the keywords\n",
    "        # for one of the grids\n",
    "        if g.dim == 2:\n",
    "            print('The assigned parameters for the 2d grid are')\n",
    "            print(d[pp.PARAMETERS][keyword_param_storage].keys())\n",
    "           \n",
    "    for e, d in gb.edges():\n",
    "        # On edges in the GridBucket, there is currently no methods for default initialization.\n",
    "\n",
    "        data = {\"normal_diffusivity\": 2e1}\n",
    "        # Add parameters: We again use keywords to identify sets of parameters.\n",
    "        d[pp.PARAMETERS] = pp.Parameters(keywords=['flow_param_edge'], dictionaries=[data])\n",
    "    \n",
    "    return gb"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Example 1\n",
    "The practical way of setting up a problem with a single variable is described here. For explanations, and hints on how to consider a more general setting, see the expanded Example 2 below. <br>\n",
    "As shown in the tutorial on single-phase flow, the equation in the mono-dimensional case is\n",
    "$$ - \\nabla \\cdot K \\nabla p = f. $$\n",
    "We expand to the mixed-dimensional version of the single-phase flow problem by solving the problem in each of the subdomains (here: fracture and matrix) and adding the flux between the subdomains\n",
    "$$ \\lambda = - \\kappa (p_{fracture} - \\texttt{tr }p_{matrix}), $$ \n",
    "with $\\kappa$ denoting the normal permeability of the fractures. For details, refer to the tutorial on single-phase flow and published papers, e.g. [this one](https://arxiv.org/abs/1802.05961).<br><br>\n",
    "We start by defining the grid bucket and assigning parameters, tagging them with a keyword. This keyword ensures that the discretizer (here tpfa, defined below) uses the right set of parameters. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The assigned parameters for the 2d grid are\n",
      "dict_keys(['source', 'mass_weight', 'second_order_tensor', 'bc', 'bc_values', 'time_step'])\n"
     ]
    }
   ],
   "source": [
    "gb, _ = pp.grid_buckets_2d.single_horizontal([2, 2], simplex=False)\n",
    "\n",
    "parameter_keyword = 'flow_param'\n",
    "gb = assign_data(gb, parameter_keyword)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, we define the variables on grids and edges and identify the individual terms of the equation we want to solve. We have an equation for the pressure on each grid (node of the GridBucket), and an equation for the mortar flux between them (edge of the bucket). The terms to be discretized are the diffusion term on the nodes ($- \\nabla \\cdot K \\nabla p$) and the coupling term $- \\kappa (p_{fracture} - \\texttt{tr }p_{matrix})$ on the edges."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the pressure variable with the same keyword on all grids\n",
    "grid_variable = 'pressure'\n",
    "# Variable name for the flux between grids, that is, the primary variable\n",
    "# on the edges in the GridBucket.\n",
    "mortar_variable = 'mortar_flux'\n",
    "\n",
    "# Identifier of the discretization operator on each grid\n",
    "operator_keyword = 'diffusion'\n",
    "# Identifier of the discretization operator between grids\n",
    "coupling_operator_keyword = 'coupling_operator'\n",
    "\n",
    "# Use a two-point flux approximation on all grids.\n",
    "# Note the keyword here: It must be the same as used when assigning the \n",
    "# parameters.\n",
    "tpfa = pp.Tpfa(parameter_keyword)\n",
    "\n",
    "# Between the grids we use a Robin type coupling (resistance to flow over a fracture).\n",
    "# Again, the keyword must be the same as used to assign data to the edge\n",
    "# The edge discretization also needs access to the corresponding discretizations \n",
    "# on the neighboring nodes\n",
    "edge_discretization = pp.RobinCoupling('flow_param_edge', tpfa, tpfa)\n",
    "\n",
    "# Loop over the nodes in the GridBucket, define primary variables and discretization schemes\n",
    "for g, d in gb:\n",
    "    # Assign primary variables on this grid. It has one degree of freedom per cell.\n",
    "    d[pp.PRIMARY_VARIABLES] = {grid_variable: {\"cells\": 1, \"faces\": 0}}\n",
    "    # Assign discretization operator for the variable.\n",
    "    # If the discretization is composed of several terms, they can be assigned\n",
    "    # by multiple entries in the inner dictionary, e.g.\n",
    "    #  {operator_keyword_1: method_1, operator_keyword_2: method_2, ...}\n",
    "    d[pp.DISCRETIZATION] = {grid_variable: {operator_keyword: tpfa}}\n",
    "    \n",
    "# Loop over the edges in the GridBucket, define primary variables and discretizations\n",
    "for e, d in gb.edges():\n",
    "    g1, g2 = gb.nodes_of_edge(e)\n",
    "    # The mortar variable has one degree of freedom per cell in the mortar grid\n",
    "    d[pp.PRIMARY_VARIABLES] = {mortar_variable: {\"cells\": 1}}\n",
    "    \n",
    "    # The coupling discretization links an edge discretization with variables\n",
    "    # and discretization operators on each neighboring grid\n",
    "    d[pp.COUPLING_DISCRETIZATION] = {\n",
    "        coupling_operator_keyword: {\n",
    "            g1: (grid_variable, operator_keyword),\n",
    "            g2: (grid_variable, operator_keyword),\n",
    "            e: (mortar_variable, edge_discretization),\n",
    "        }\n",
    "    }\n",
    "    d[pp.DISCRETIZATION_MATRICES] = {'flow_param_edge': {}}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The task of assembling the linear system is left to a dedicated object, called an `Assembler`, whih again relies on a `DofManager` to keep track of the ordering of unknowns (for more, see below).\n",
    "\n",
    "Discretization and assembly of the global linear system can in this case be carried out by a single function call. Note that for some problems, notably poro-elasticity, this is not possible, then discretization must be carried out first.\n",
    "\n",
    "Below, A is the global linear system, and b is the corresponding right hand side, and we obtain the pressure solution by solving the system."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "dof_manager = pp.DofManager(gb)\n",
    "\n",
    "assembler = pp.Assembler(gb, dof_manager)\n",
    "assembler.discretize()\n",
    "\n",
    "# Assemble the linear system, using the information stored in the GridBucket\n",
    "A, b = assembler.assemble_matrix_rhs()\n",
    "\n",
    "pressure = sps.linalg.spsolve(A, b)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The parameters assigned above will not yield a well-posed problem, thus the solve will likely produce a warning about the matrix being singular. This can be ignored. <br>\n",
    "\n",
    "The ordering of the unknowns in the global linear system will vary depending on how the components in the GridBucket and the unknowns are traversed. Untangling the ordering is a two-stage process:\n",
    "1. The system is ordered as a block system, with one block per combination of primary variable and grid or edge (between grids). This information is stored in the attribute `dof_manager.block_dof`.\n",
    "2. For each block, the local degrees of freedom can be obtained from the attribute `dof_manager.full_dof`.\n",
    "\n",
    "To get the block number of a specific primary variable, we need the identifier of the relevant component in the GridBucket (either the grid, or the edge between grids), and the variable name."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[ 50. -10.   0.   0.]\n",
      " [-10.  50.   0.   0.]\n",
      " [  0.   0.  50. -10.]\n",
      " [  0.   0. -10.  50.]]\n"
     ]
    }
   ],
   "source": [
    "# Getting the grids is easy, there is one in each dimension\n",
    "g_2d = gb.grids_of_dimension(2)[0]\n",
    "g_1d = gb.grids_of_dimension(1)[0]\n",
    "\n",
    "# Formally loop over the edges, there is a single one\n",
    "for e, _ in gb.edges():\n",
    "    continue\n",
    "\n",
    "# Now, the block ordering is obtained, for the 2d grid as \n",
    "block_2d = dof_manager.block_dof[(g_2d, grid_variable)]\n",
    "\n",
    "# full_dof contains the number of dofs per block. To get a global ordering, use\n",
    "global_dof = np.r_[0, np.cumsum(dof_manager.full_dof)]\n",
    "\n",
    "# Get 2d dofs\n",
    "global_dof_2d = np.arange(global_dof[block_2d], global_dof[block_2d+1])\n",
    "# Print the relevant part of the system matrix\n",
    "print(A.toarray()[global_dof_2d, :][:, global_dof_2d])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Example 2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The first example showed how to work with the assembler in reletively simple cases. In this second example, we aim to illustrate the full scope of the assembler, including:\n",
    "* General assignment of variables on different grid components (fracture, matrix, etc.):\n",
    "    * Different number of variables on each grid component\n",
    "    * Different names for variables (a relevant case could be to use 'temperature' on one domain, 'enthalpy' on another, with an appropriate coupling)\n",
    "* General coupling schemes between different grid components:\n",
    "    * Multiple coupling variables\n",
    "    * Couplings related to different variables and discretization schemes on the neighboring grids.\n",
    "* Multiple discretization operators applied to the same term / equation on different grid components\n",
    "\n",
    "\n",
    "The example that incorporates all these features are necessarily quite complex and heavy on notation. As such it should be considered as a reference for how to use the functionality, more than a simulation of any real physical system.\n",
    "\n",
    "We define two primary variables on the nodes and three coupling variables. The resulting system will be somewhat arbitrary, in that it may not reflect any standard physics, but it should better illustrate what is needed for a multi-physics problem.\n",
    "\n",
    "First we extend the data assignment method."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "def assign_data_2(gb, keyword_param_storage, keyword_param_storage_2=None):\n",
    "    # Method to assign data.\n",
    "    for g, d in gb:\n",
    "        # This keyword is used to define which set of default parameters to pick\n",
    "        # Replace with 'transport' or 'mechanics' if needed\n",
    "        default_parameter_type = 'flow'  \n",
    "\n",
    "        # Assign a non-default permeability, for illustrative purposes\n",
    "        if g.dim == 2:\n",
    "            kxx = 10 * np.ones(g.num_cells)\n",
    "        else:\n",
    "            kxx = 0.1 * np.ones(g.num_cells)\n",
    "\n",
    "        perm = pp.SecondOrderTensor(kxx)\n",
    "\n",
    "        # Create a dictionary to override the default parameters\n",
    "        specified_parameters = {'second_order_tensor': perm}\n",
    "\n",
    "        #\n",
    "\n",
    "        # Define the \n",
    "        pp.initialize_default_data(g, d, default_parameter_type, specified_parameters,\n",
    "                                   keyword_param_storage)\n",
    "\n",
    "        # Internally to the Parameter class, the parameters are stored as dictionaries.\n",
    "        # To illustrate how to access specific sets of parameters, print the keywords\n",
    "        # for one of the grids\n",
    "        if g.dim == 2 and not keyword_param_storage_2:\n",
    "            print('The assigned parameters for the 2d grid are')\n",
    "            print(d[pp.PARAMETERS][keyword_param_storage].keys())\n",
    "           \n",
    "        # For one example below, we will need two different parameter sets.\n",
    "        # Define a second set, with default values only.\n",
    "        if keyword_param_storage_2:\n",
    "            pp.initialize_default_data(g, d, default_parameter_type, keyword = keyword_param_storage_2)\n",
    "        \n",
    "\n",
    "    for e, d in gb.edges():\n",
    "        # On edges in the GridBucket, there is currently no methods for default initialization.\n",
    "\n",
    "        data = {\"normal_diffusivity\": 2e1}\n",
    "        # Add parameters: We again use keywords to identify sets of parameters.\n",
    "        if keyword_param_storage_2 is not None:\n",
    "            # There are actually three parameters here ('two_parameter_sets' refers to the nodes)\n",
    "            # since we plan on using in total three mortar variables in this case\n",
    "            d[pp.PARAMETERS] = pp.Parameters(keywords=['flow_param_edge',\n",
    "                                                       'second_flow_param_edge',\n",
    "                                                       'third_flow_param_edge'],\n",
    "                                             dictionaries=[data, data, data])\n",
    "        else:\n",
    "            d[pp.PARAMETERS] = pp.Parameters(keywords=['flow_param_edge'], dictionaries=[data])\n",
    "    \n",
    "    return gb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define a grid\n",
    "gb, _ = pp.grid_buckets_2d.single_horizontal([4, 4], simplex=False)\n",
    "parameter_keyword = 'flow_param'\n",
    "parameter_keyword_2 = 'second_flow_param'\n",
    "gb = assign_data_2(gb, parameter_keyword, parameter_keyword_2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Primary variables must be defined on each component of the GridBucket.\n",
    "\n",
    "On the first grid we use a cell centered method which has one primary variable \"pressure\". \n",
    "On the second grid, we use a mixed method with both pressure and fluxes combined into one primary variable.\n",
    "\n",
    "The temperature is tagged with the same keyword on both grids."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Variable keywords first grid\n",
    "grid_1_pressure_variable = 'pressure'\n",
    "grid_1_temperature_variable = 'temperature'\n",
    "# Variable keywords second grid\n",
    "grid_2_pressure_variable = 'flux_pressure'\n",
    "grid_2_temperature_variable = 'temperature'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next we assign a keyword to the coupling terms between the grid. We will have three coupling variables;\n",
    "one for the fluid flux, and one for each of the diffusive terms in the temperature equation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Coupling variable for pressure\n",
    "mortar_variable_pressure = 'mortar_flux_pressure'\n",
    "# Coupling variable for advective temperature flux\n",
    "mortar_variable_temperature_1 = 'mortar_flux_diffusion'\n",
    "mortar_variable_temperature_2 = 'mortar_flux_diffusion_2'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We now give a keyword to the operators."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Identifier of the discretization operator for pressure discretizaiton\n",
    "operator_keyword_pressure = 'pressure_diffusion'\n",
    "\n",
    "# identifier for the temperature discretizations.\n",
    "# THIS IS WEIRD: The intention is to illustrate the use of two discretization operators for \n",
    "# a single variable. The natural option in this setting is advection-diffusion, but that\n",
    "# requires either the existence of a Darcy flux, or tighter coupling with the pressure equation.\n",
    "# Purely for illustrative purposes, we instead use a double diffusion model. There you go.\n",
    "operator_keyword_temperature_1 = 'diffusion'\n",
    "operator_keyword_temperature_2 = 'diffusion_2'\n",
    "\n",
    "# Identifier of the discretization operator between grids\n",
    "coupling_pressure_keyword = 'coupling_operator_pressure'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "So far we have only defined the keywords needed for the discretizations to obtain the correct parameters\n",
    "and couplings. Next, we create the discretization objects"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Pressure diffusion discretization\n",
    "tpfa_flow = pp.Tpfa(parameter_keyword)\n",
    "vem_flow = pp.MVEM(parameter_keyword)\n",
    "# Temperature diffusion discretization\n",
    "tpfa_temperature = pp.Tpfa(parameter_keyword_2)\n",
    "mpfa_temperature = pp.Mpfa(parameter_keyword_2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Discretization operators on the coupling conditions, chosen to illustrate the framework.\n",
    "Note that in all cases, the coupling conditions need a separate keyword, which should \n",
    "correspond to an assigned set of data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "# One term couples two pressure / flow variables\n",
    "edge_discretization_flow = pp.RobinCoupling('flow_param_edge', tpfa_flow, vem_flow)\n",
    "\n",
    "# The second coupling is of mpfa on one domain, and tpfa on the other, both for temperature\n",
    "edge_discretization_temperature_diffusion_1 = pp.RobinCoupling('second_flow_param_edge',\n",
    "                                                               mpfa_temperature, tpfa_temperature)\n",
    "\n",
    "# The third coupling is of tpfa for flow with mpfa for temperature\n",
    "edge_discretization_temperature_diffusion_2 = pp.RobinCoupling('third_flow_param_edge',\n",
    "                                                               tpfa_flow, mpfa_temperature)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Loop over the nodes in the GridBucket, define primary variables and discretization schemes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "for g, d in gb:\n",
    "    # Assign primary variables on this grid. \n",
    "    if g.dim == 2:\n",
    "        # Both pressure and temperature are represented as cell centered variables\n",
    "        d[pp.PRIMARY_VARIABLES] = {grid_1_pressure_variable: {\"cells\": 1, \"faces\": 0},\n",
    "                                   grid_1_temperature_variable: {\"cells\": 1}}\n",
    "\n",
    "        # The structure of the discretization assignment is: For each variable, give a \n",
    "        # pair of operetor identifications (usually a string) and a discretizaiton method.\n",
    "        # If a variable is identified with several discretizations, say, advection and diffusion,\n",
    "        # several pairs can be assigned.\n",
    "        \n",
    "        # For pressure, use tpfa.\n",
    "        # For temperature, use two discretizations, respectively tpfa and mpfa\n",
    "        d[pp.DISCRETIZATION] = {grid_1_pressure_variable: {operator_keyword_pressure: tpfa_flow},\n",
    "                                grid_1_temperature_variable: {operator_keyword_temperature_1: tpfa_temperature,\n",
    "                                                              operator_keyword_temperature_2: mpfa_temperature}}\n",
    "    else:  #g.dim == 1\n",
    "        # Pressure is discretized with flux-pressure combination, temperature with cell centered variables\n",
    "        d[pp.PRIMARY_VARIABLES] = {grid_2_pressure_variable: {\"cells\": 1, \"faces\": 1},\n",
    "                                   grid_2_temperature_variable: {\"cells\": 1}}\n",
    "        # For pressure, use vem.\n",
    "        # For temperature, only discretize once, with tpfa\n",
    "        d[pp.DISCRETIZATION] = {grid_2_pressure_variable: {operator_keyword_pressure: vem_flow},\n",
    "                                grid_2_temperature_variable: {operator_keyword_temperature_1: tpfa_temperature}}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Loop over the edges in the GridBucket, define primary variables and discretizations.\n",
    "\n",
    "Notice how coupling discretizations are assigned as a dictionary, one per coupling term on each edge. For each term, the coupling contains an inner dictionary, with the keys being the edge and the two neighboring grids. For the edge, the values are the name of the mortar variable, and the discretization object to be applied. For the grids, the values are the variable name on the grid, and the keyword identifying the discretization operator, as specified in the loop over nodes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "for e, d in gb.edges():\n",
    "    # \n",
    "    g1, g2 = gb.nodes_of_edge(e)\n",
    "    \n",
    "    # The syntax used in the problem setup assumes that g1 has dimension 2\n",
    "    if g1.dim < g2.dim:\n",
    "        g2, g1 = g1, g2\n",
    "                                \n",
    "    # The mortar variable has one degree of freedom per cell in the mortar grid\n",
    "    d[pp.PRIMARY_VARIABLES] = {mortar_variable_pressure: {\"cells\": 1},\n",
    "                               mortar_variable_temperature_1: {\"cells\": 1},\n",
    "                               mortar_variable_temperature_2: {\"cells\": 1},\n",
    "                              }\n",
    "    \n",
    "    # Coupling discretizations\n",
    "    d[pp.COUPLING_DISCRETIZATION] = {\n",
    "        # The flow discretization couples tpfa on one domain with vem on the other\n",
    "        'edge_discretization_flow': {\n",
    "            g1: (grid_1_pressure_variable, operator_keyword_pressure),\n",
    "            g2: (grid_2_pressure_variable, operator_keyword_pressure),\n",
    "            e: (mortar_variable_pressure, edge_discretization_flow),\n",
    "        },\n",
    "        # The first temperature mortar couples one of the temperature discretizations on grid 1\n",
    "        # with the single tempearture discretization on the second grid\n",
    "        # As a side remark, the keys in the outer dictionary are never used, except from debugging,\n",
    "        # but a dictionary seemed a more natural option than a list.\n",
    "        'the_keywords_in_this_dictionary_can_have_any_value': {  \n",
    "            g1: (grid_1_temperature_variable, operator_keyword_temperature_2),\n",
    "            g2: (grid_2_temperature_variable, operator_keyword_temperature_1),\n",
    "            e: (mortar_variable_temperature_1, edge_discretization_temperature_diffusion_1),\n",
    "        },   \n",
    "        # Finally, the third coupling\n",
    "        'second_edge_discretization_temperature': {\n",
    "            # grid_1_variable_1 gives pressure variable, then identify the discretization object\n",
    "            g1: (grid_1_pressure_variable, operator_keyword_pressure),\n",
    "            # grid_2_variable_2 gives temperature, then use the keyword that was used to identify mpfa\n",
    "            # (and not the one for tpfa, would have been operator_keyword_temperature_1)\n",
    "            g2: (grid_2_temperature_variable, operator_keyword_temperature_2),\n",
    "            e: (mortar_variable_temperature_2, edge_discretization_temperature_diffusion_2),\n",
    "        }       \n",
    "        \n",
    "    }\n",
    "    d[pp.DISCRETIZATION_MATRICES] = {'flow_param_edge': {},\n",
    "                                     'second_flow_param_edge': {},\n",
    "                                     'third_flow_param_edge': {}\n",
    "                                    }"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We have now assigned all the data. The task of assembling the linear system is left to a dedicated object:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "dof_manager = pp.DofManager(gb)\n",
    "assembler = pp.Assembler(gb, dof_manager)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Discretization and assembly of the global linear system can again be carried out by separate function calls."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[0. 0. 0. 0. 0. 0. 0. 0.]\n",
      " [0. 0. 0. 0. 0. 0. 0. 0.]\n",
      " [0. 0. 0. 0. 0. 0. 0. 0.]\n",
      " [0. 0. 0. 0. 0. 0. 0. 0.]\n",
      " [0. 0. 0. 0. 0. 0. 0. 1.]\n",
      " [0. 0. 0. 0. 0. 0. 1. 0.]\n",
      " [0. 0. 0. 0. 0. 1. 0. 0.]\n",
      " [0. 0. 0. 0. 1. 0. 0. 0.]\n",
      " [0. 0. 0. 1. 0. 0. 0. 0.]\n",
      " [0. 0. 1. 0. 0. 0. 0. 0.]\n",
      " [0. 1. 0. 0. 0. 0. 0. 0.]\n",
      " [1. 0. 0. 0. 0. 0. 0. 0.]\n",
      " [0. 0. 0. 0. 0. 0. 0. 0.]\n",
      " [0. 0. 0. 0. 0. 0. 0. 0.]\n",
      " [0. 0. 0. 0. 0. 0. 0. 0.]\n",
      " [0. 0. 0. 0. 0. 0. 0. 0.]]\n"
     ]
    }
   ],
   "source": [
    "# Discretize, then Assemble the linear system, using the information stored in the GridBucket\n",
    "assembler.discretize()\n",
    "A, b = assembler.assemble_matrix_rhs()\n",
    "\n",
    "# Pick out part of the discretization associated with the third mortar variable\n",
    "g_2d = gb.grids_of_dimension(2)[0]\n",
    "\n",
    "# Formally loop over the edges, there is a single one\n",
    "for e, _ in gb.edges():\n",
    "    continue\n",
    "\n",
    "# Now, the block ordering is obtained, for the 2d grid as \n",
    "block_2d_pressure = dof_manager.block_dof[(g_2d, grid_1_pressure_variable)]\n",
    "block_e_third_mortar = dof_manager.block_dof[(e, mortar_variable_temperature_2)]\n",
    "\n",
    "# full_dof contains the number of dofs per block. To get a global ordering, use\n",
    "global_dof = np.r_[0, np.cumsum(dof_manager.full_dof)]\n",
    "\n",
    "# Get 2d dofs\n",
    "global_dof_2d_pressure = np.arange(global_dof[block_2d_pressure], global_dof[block_2d_pressure+1])\n",
    "global_dof_e_temperature = np.arange(global_dof[block_e_third_mortar], global_dof[block_e_third_mortar+1])\n",
    "# Print the relevant part of the system matrix\n",
    "print(A.toarray()[global_dof_2d_pressure, :][:, global_dof_e_temperature])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
